from mealpy import SCA, FloatVar, GA, PSO, SCSO, COA, WHO, BBOA, ACOR
from SAGIN_Progrem.ENV import Env
import matplotlib.pyplot as plt
from copy import deepcopy


# PSO.P_PSO 参考文献2019年

def run(e, multi_action):
    off_sa = [[], [], [], [], []]
    off_sa_list = [[], [], [], [], []]
    sa_com = []
    # 卫星要本地计算的任务：无人机卸载的+收集的
    for task in e.sate[0].collect(e.task_list_passion_sa[e.now_step], e.step_time):
        sa_com.append(task)
    for i in range(0, e.uav_num):
        arrive_time_list = e.uav_list[i].collect_c(e.task_list_passion_uav[i][e.now_step], e.step_time)
        # 在计算队列中计算要卸载的任务
        off_sa_temp = e.uav_list[i].compute_task_n2_(arrive_time_list, e.task_list_passion_uav[i][e.now_step],
                                                        e.step_time, multi_action[i], e.now_step + 1)
        # 获取所有卸载到卫星的任务
        off_sa[i].append(off_sa_temp)
        # 获取能传输到卫星的任务
        off_sa_list_temp = e.uav_list[i].off_to_sate_c(off_sa[i][0], e.step_time)  # 卸载到卫星的
        # 统计
        off_sa_list[i].append(off_sa_list_temp)
        for task in off_sa_list[i][0]:
            sa_com.append(task)
    # 卫星计算
    e.sa_com = sa_com
    off_cloud_temp = e.sate[0].compute_task_(sa_com, e.step_time, multi_action[e.uav_num],
                                             e.now_step + 1)
    off_cloud = e.sate[0].off_to_cloud(off_cloud_temp, e.step_time)
    # 云计算
    e.cloud[0].compute_task(off_cloud, e.step_time, e.now_step + 1)

def run_(e, multi_action):
    off_sa = [[], [], [], [], []]
    off_sa_list = [[], [], [], [], []]
    sa_com = []
    # 卫星要本地计算的任务：无人机卸载的+收集的
    for task in e.sate[0].collect(e.task_list_passion_sa, e.step_time):
        sa_com.append(task)
    for i in range(0, e.uav_num):
        arrive_time_list = e.uav_list[i].collect_c(e.task_list_passion_uav[i], e.step_time)
        # 在计算队列中计算要卸载的任务
        off_sa_temp = e.uav_list[i].compute_task_n2_(arrive_time_list, e.task_list_passion_uav[i],
                                                        e.step_time, multi_action[i], e.now_step + 1)
        # 获取所有卸载到卫星的任务
        off_sa[i].append(off_sa_temp)
        # 获取能传输到卫星的任务
        off_sa_list_temp = e.uav_list[i].off_to_sate_c(off_sa[i][0], e.step_time)  # 卸载到卫星的
        # 统计
        off_sa_list[i].append(off_sa_list_temp)
        for task in off_sa_list[i][0]:
            sa_com.append(task)
    # 卫星计算
    e.sa_com = sa_com
    off_cloud_temp = e.sate[0].compute_task_(sa_com, e.step_time, multi_action[e.uav_num],
                                             e.now_step + 1)
    off_cloud = e.sate[0].off_to_cloud(off_cloud_temp, e.step_time)
    # 云计算
    e.cloud[0].compute_task(off_cloud, e.step_time, e.now_step + 1)

def step(multi_actions):
    env_temp = Env(6, MAX_STEP)
    env_temp.app_list = deepcopy(env.app_list)
    env_temp.uav_list = deepcopy(env.uav_list)
    env_temp.sate = deepcopy(env.sate)
    env_temp.cloud = deepcopy(env.cloud)
    for i in range(env_temp.uav_num):
        env_temp.task_list_passion_uav[i] = deepcopy(env.task_list_passion_uav[i][env.now_step])
    env_temp.task_list_passion_sa = deepcopy(env.task_list_passion_sa[env.now_step])
    env_temp.now_step = env.now_step
    multi_rewards = []
    # 首先获取上轮在各队列中剩余的任务
    uav_last_task = []
    sa_last_task = []
    for i in range(env_temp.uav_num):
        temp = []
        for j in env_temp.uav_list[i].computing_queue.items:
            temp.append(j)
        for j in env_temp.uav_list[i].offloading_sa_queue.items:
            temp.append(j)
        uav_last_task.append(temp)
    for i in env_temp.sate[0].computing_queue.items:
        sa_last_task.append(i)
    for i in env_temp.sate[0].offloading_queue.items:
        sa_last_task.append(i)
    run_(env_temp, multi_actions)
    # 更新所有agent的奖励值
    # reward等于完成率乘以负的平均延迟
    finish_rate = []  # 完成率  一轮五个无人机分别的完成率
    sa_finish = 0  # 卫星收集的个数
    uav_finish = []  # 计算一共完成了多少任务
    for state_i in range(env_temp.uav_num):
        uav_finish_num = 1  # 计算一共完成了多少任务
        delay = 0  # 总延迟
        delay_t = 0
        delay_c = 0
        delay_w = 0
        for task in env_temp.task_list_passion_uav[state_i]:
            if not task.finish == 0:
                uav_finish.append(task)
                uav_finish_num += 1
                delay_t += task.delay['transmission']
                delay_c += task.delay['compute']
                delay_w += task.delay['wait']
                delay += delay_t + delay_c + delay_w
                delay_t_ = task.delay['transmission']
                delay_c_ = task.delay['compute']
                delay_w_ = task.delay['wait']
                delay_ = delay_t_ + delay_c_ + delay_w_
        for task in uav_last_task[state_i]:
            if not task.finish == 0:
                uav_finish.append(task)
                uav_finish_num += 1
                delay_t += task.delay['transmission']
                delay_c += task.delay['compute']
                delay_w += task.delay['wait']
                delay += delay_t + delay_c + delay_w
                delay_t_ = task.delay['transmission']
                delay_c_ = task.delay['compute']
                delay_w_ = task.delay['wait']
                delay_ = delay_t_ + delay_c_ + delay_w_
        delay = delay_t + delay_c + delay_w
        finish_rate.append(
            uav_finish_num / (len(env_temp.task_list_passion_uav[state_i]) + len(uav_last_task[state_i])))
        k = (7 ** (1 - finish_rate[state_i])) - 0.7
        reward = 130 * finish_rate[state_i] - k * delay / uav_finish_num
        multi_rewards.append(reward)
    sa_delay = 0
    for task in env_temp.sa_com:
        if not task.finish == 0:
            sa_finish += 1
            sa_delay += task.delay['wait'] + task.delay['compute'] + task.delay['transmission']
    for task in sa_last_task:
        if not task.finish == 0:
            sa_finish += 1
            sa_delay += task.delay['wait'] + task.delay['compute'] + task.delay['transmission']
    finish_rate = sa_finish / (len(env_temp.sa_com) + len(sa_last_task))
    k = (7 ** (1 - finish_rate)) - 0.7
    try:
        reward_sa = 100 * finish_rate - k * sa_delay / sa_finish
    except:
        reward_sa = 100 * finish_rate - 130
    multi_rewards.append(reward_sa)
    return sum(multi_rewards)


MAX_STEP = 25

if __name__ == '__main__':

    env = Env(6, MAX_STEP)
    env.init_entity()
    problem_dict = {
        "obj_func": step,
        "bounds": FloatVar(lb=[0, 0, 0, 0, 0, 0], ub=[1, 1, 1, 1, 1, 1]),
        "minmax": "max",
    }
    # .........................COA...........................
    COA_global_best_fit_list = []
    COA_current_best_fit_list = []
    env.rander()
    for i in range(MAX_STEP):
        env.now_step = i
        optimizer_COA = COA.OriginalCOA(epoch=50, pop_size=100, n_coyotes=5)
        optimizer_COA.solve(problem_dict)
        run(env, optimizer_COA.g_best.solution)
        COA_global_best_fit_list.append(optimizer_COA.history.list_global_best_fit)
        COA_current_best_fit_list.append(optimizer_COA.history.list_current_best_fit)
    COA_total_finish = []
    COA_sa, COA_sa_, COA_uav, COA_uav_ = 0, 0, 0, 0
    for i in range(MAX_STEP):
        for task in env.task_list_passion_sa[i]:
            COA_sa_ += 1
            if not task.finish == 0:
                COA_sa += 1
                COA_total_finish.append(task)
        for j in range(env.agents_num - 1):
            for task in env.task_list_passion_uav[j][i]:
                COA_uav_ += 1
                if not task.finish == 0:
                    COA_uav += 1
                    COA_total_finish.append(task)
    COA_finish_rate = (COA_sa + COA_uav) / (COA_sa_ + COA_uav_)
    COA_delay_list = []
    for task in COA_total_finish:
        delay = task.delay['wait'] + task.delay['compute'] + task.delay['transmission']
        COA_delay_list.append(delay)